import json
from docx import Document
from docx.shared import Pt
from docx.oxml.ns import qn
from typing import List, Dict

import streamlit as st

import gene_tst
from logs import logger as log
from langchain.memory import ConversationBufferWindowMemory
import sys
from service import bid_service,gen_catalog


print(f'sys.path={sys.path}')

# """
# 用streamlit帮我们做一个简单的网页可以用来调用我们的机器人问答
# """
st.title('💡招标信息解析 💡')  # 用streamlit app创建一个标题
# 使用streamlit.message_state存储记忆(memory)和状态信息(messages)
if 'memory' not in st.session_state:
    # 暂时为了节省token，先设置k=1
    st.session_state['memory'] = ConversationBufferWindowMemory(return_messages=True, k=1)
    st.session_state['messages'] = [{'role': 'ai', 'content': '我是一名合格的招投标助手，请问有什么可以帮助您的？'}]

# 页面渲染聊天记录
for message in st.session_state['messages']:
    st.chat_message(message['role']).write(message['content'], unsafe_allow_html=True)

# 获取聊天框输入的内容
query = st.chat_input()

class GenerateChapter:
    """
    文章生成器
    """
    # 将 method_to_chinese_option 转化为静态变量
    method_to_chinese_option = {
        "agent_route_generation": "Agent路线生成",
        "database_matching": "数据库匹配",
        "tender_document_parsing": "招标书解析生成",
    }

    def __init__(self, chapters: Dict, font_size: int = 16, font_name: str = u'仿宋_GB2312', doc_styles: str = 'Normal') -> None:
        """
        初始化生成器，需要传入章节名

        :param chapters: 章节名的列表
        :param font_size: 字体大小
        :param font_name: 字体
        :param doc_styles: 文本编辑格式
        """
        self.chapters = chapters
        self.font_size = font_size
        self.font_name = font_name
        self.doc_styles = doc_styles

        self.init_state()



    def init_state(self) -> None:
        log.info("init_state exec...")
        # 初始化 session_state
        if "generated_sections" not in st.session_state:
            log.info("init_state generated_sections exec...")
            st.session_state.generated_sections = set()  # 存储已生成的章节、子小节和子子小节标题
        if "chapter_docs" not in st.session_state:
            log.info("init_state chapter_docs exec...")
            st.session_state.chapter_docs = {}  # 存储每个章节的文档对象
        if "selected_chapter" not in st.session_state:
            log.info("init_state selected_chapter exec...")
            st.session_state.selected_chapter = None  # 存储当前选中的章
        if "selected_section" not in st.session_state:
            log.info("init_state selected_section exec...")
            st.session_state.selected_section = None  # 存储当前选中的小节
        if "selected_subsection" not in st.session_state:
            log.info("init_state selected_subsection exec...")
            st.session_state.selected_subsection = None  # 存储当前选中的子节
        if "selected_subsubsection" not in st.session_state:
            log.info("init_state selected_subsubsection exec...")
            st.session_state.selected_subsubsection = None  # 存储当前选中的子子节
        if "selected_method" not in st.session_state:
            log.info("init_state selected_method exec...")
            st.session_state.selected_method = None  # 存储当前选中的生成方法



    def begin(self, chapter_title: str) -> Document:
        """
        生成章节文档

        :param chapter_title: 章的名称
        :return: 有关章的 docx 文件
        """
        doc = Document()
        doc.styles[self.doc_styles].font.size = Pt(self.font_size)
        doc.styles[self.doc_styles].font.name = self.font_name
        doc.styles[self.doc_styles].element.rPr.rFonts.set(qn('w:eastAsia'), self.font_name)
        doc.add_heading(chapter_title, level=1)  # 添加章节标题

        return doc


    def get_content_from_db(self, section_title):
        """
        数据库取数据
        """
        pass


    def save_document(self, doc: Document, chapter_title: str) -> None:
        """
        保存文档

        :param doc: 文档
        :param chapter_title: 章名
        """
        file_all_name = chapter_title
        file_name = f'{file_all_name}.docx'
        # 保存章节文档
        doc.save(file_name)
        st.success(f"{chapter_title} 章节生成成功，文件名为：{file_name}")


    def generate_content(self, method: str, section_title: str) -> str:
        """
        根据选择的方法生成内容

        :param method: 生成方法（Agent路线生成、数据库匹配、招标书解析生成）
        :param section_title: 小节标题
        :return: 生成的内容
        """
        if method == "agent_route_generation":
            return f"这是通过 Agent路线生成 生成的内容，标题为：{section_title}。"
        elif method == "database_matching":
            return f"这是通过 数据库匹配 生成的内容，标题为：{section_title}。"
        elif method == "tender_document_parsing":
            # content = bid_service.chat(section_title)
            return f"这是通过 招标书解析生成 生成的内容，标题为：{section_title}。"


    def generate_subsection(self, doc: Document, subsection_title: str, method: str) -> None:
        """
        生成子小节

        :param doc: 文档
        :param subsection_title: 子小节名
        :param method: 生成方法
        """

        if subsection_title in st.session_state.generated_sections:
            st.warning(f"{subsection_title} 已经生成过了！")

            return

        st.session_state.generated_sections.add(subsection_title)

        # 添加子小节标题
        heading = doc.add_heading(subsection_title, level=3)
        heading_run = heading.runs[0]
        heading_run.font.size = Pt(self.font_size - 4)
        heading_run.font.name = self.font_name
        heading_run.element.rPr.rFonts.set(qn('w:eastAsia'), self.font_name)  # 设置中文字体

        # 根据选择的方法生成内容
        content = self.generate_content(method, subsection_title)

        paragraph = doc.add_paragraph()
        paragraph.paragraph_format.first_line_indent = 0
        paragraph.paragraph_format.element.pPr.ind.set(qn("w:firstLineChars"), '200')
        paragraph.add_run(content).font.size = Pt(self.font_size - 4)

        return


    def generate_subsubsection(self, doc: Document, subsubsection_title: str, method: str) -> None:
        """
        生成子子小节

        :param doc: 文档
        :param subsubsection_title: 子子节名
        :param method: 生成方法
        """

        if subsubsection_title in st.session_state.generated_sections:
            st.warning(f"{subsubsection_title} 已经生成过了！")
            return

        st.session_state.generated_sections.add(subsubsection_title)

        # 添加子子小节标题
        heading = doc.add_heading(subsubsection_title, level=3)
        heading_run = heading.runs[0]
        heading_run.font.size = Pt(self.font_size - 3)
        heading_run.font.name = self.font_name
        heading_run.element.rPr.rFonts.set(qn('w:eastAsia'), self.font_name)  # 设置中文字体

        # 根据选择的方法生成内容
        content = self.generate_content(method, subsubsection_title)

        paragraph = doc.add_paragraph()
        paragraph.paragraph_format.first_line_indent = 0
        paragraph.paragraph_format.element.pPr.ind.set(qn("w:firstLineChars"), '200')
        paragraph.add_run(content).font.size = Pt(self.font_size - 4)

        return


    def generate_section(self, doc: Document, section_title: str, method: str) -> None:
        """
        生成小节

        :param doc: 文档
        :param section_title: 小节名
        :param subsections: 子小节列表
        :param method: 生成方法
        """

        if section_title in st.session_state.generated_sections:
            st.warning(f"{section_title} 已经生成过，无法重复生成。")

            return

        st.session_state.generated_sections.add(section_title)
        # 添加小节标题
        heading = doc.add_heading(section_title, level=2)
        heading_run = heading.runs[0]
        heading_run.font.size = Pt(self.font_size - 3)
        heading_run.font.name = self.font_name
        heading_run.element.rPr.rFonts.set(qn('w:eastAsia'), self.font_name)  # 设置中文字体

        # content = '梨树县，隶属于吉林省四平市，位于吉林省西南部，地处松辽平原腹地。梨树县东与长春市隔东辽河相望，南与四平市区接壤，西与辽宁省铁岭市毗连，北与双辽市以东辽河为界。介于北纬43°02′-43°46′，东经123°45′～124°53′之间，总面积3232 平方千米。截至2022 年10 月，梨树县下辖3 个街道、15 个镇、6 个乡。'
        content = self.generate_content(method, section_title)

        paragraph = doc.add_paragraph()
        paragraph.paragraph_format.first_line_indent = 0
        paragraph.paragraph_format.element.pPr.ind.set(qn("w:firstLineChars"), '200')
        paragraph.add_run(content).font.size = Pt(self.font_size - 4)

        return


    def create_launcher(self) -> None:
        """
        创建启动器窗口
        """
        with st.sidebar:
            st.title("目录")

            # 显示章的生成按钮
            for chapter in self.chapters:
                chapter_title = chapter["title"]
                if st.button(f"生成 {chapter_title}"):
                    st.session_state.selected_chapter = chapter_title
                    st.session_state.selected_section = None  # 重置小节
                    st.session_state.selected_subsection = None  # 重置子节
                    st.session_state.selected_subsubsection = None  # 重置子子节
                    st.session_state.selected_method = None  # 重置生成方法

            # 如果选择了章，显示该章的所有小节、子节和子子节按钮
            if st.session_state.selected_chapter:
                st.header(f"生成 {st.session_state.selected_chapter}")
                chapter_data = next(ch for ch in self.chapters if ch["title"] == st.session_state.selected_chapter)

                # 显示小节按钮
                for section in chapter_data.get("sections", []):
                    section_title = section["title"]
                    if st.button(f"生成 {section_title}"):
                        st.session_state.selected_section = section_title
                        st.session_state.selected_subsection = None  # 重置子节
                        st.session_state.selected_subsubsection = None  # 重置子子节
                        st.session_state.selected_method = None  # 重置生成方法

                    # 显示子节按钮
                    for subsection in section.get("sections", []):
                        subsection_title = subsection["title"]
                        if st.button(f"生成 {subsection_title}"):
                            st.session_state.selected_subsection = subsection_title
                            st.session_state.selected_subsubsection = None  # 重置子子节
                            st.session_state.selected_method = None  # 重置生成方法

                        # 显示子子节按钮
                        for subsubsection in subsection.get("sections", []):
                            subsubsection_title = subsubsection["title"]
                            if st.button(f"生成 {subsubsection_title}"):
                                st.session_state.selected_subsubsection = subsubsection_title
                                st.session_state.selected_method = None  # 重置生成方法

            # 如果选择了小节、子节或子子节，显示生成方法选择框
            if st.session_state.selected_section or st.session_state.selected_subsection or st.session_state.selected_subsubsection:
                selected_title = (
                        st.session_state.selected_subsubsection or
                        st.session_state.selected_subsection or
                        st.session_state.selected_section)
                st.subheader(f"生成 {selected_title}")
                # 修改: 将中文选项映射到英文键
                method_chinese_options = list(GenerateChapter.method_to_chinese_option.values())
                method_english_keys = list(GenerateChapter.method_to_chinese_option.keys())
                method_index = st.radio(
                    "选择生成方法：",
                    method_chinese_options,
                    key="method_selector")
                method = method_english_keys[method_chinese_options.index(method_index)]
                st.session_state.selected_method = method

                # 确认生成按钮
                if st.button("确认生成"):
                    chapter_title = st.session_state.selected_chapter
                    selected_title = (
                            st.session_state.selected_subsubsection or
                            st.session_state.selected_subsection or
                            st.session_state.selected_section)

                    if chapter_title not in st.session_state.chapter_docs:
                        # 创建章节文档并保存到 session_state
                        doc = self.begin(chapter_title)
                        st.session_state.chapter_docs[chapter_title] = doc
                    else:
                        # 从 session_state 中获取已有的文档
                        doc = st.session_state.chapter_docs[chapter_title]

                    # 根据选择的标题生成内容
                    if st.session_state.selected_section:
                        self.generate_section(doc, selected_title, method)
                    elif st.session_state.selected_subsection:
                        self.generate_subsection(doc, selected_title, method)
                    elif st.session_state.selected_subsubsection:
                        self.generate_subsubsection(doc, selected_title, method)

                    st.success(f"{selected_title} 生成完成！")

            # 保存文档按钮
            if st.session_state.selected_chapter and st.button(f"保存 {st.session_state.selected_chapter}"):
                if st.session_state.selected_chapter in st.session_state.chapter_docs:
                    self.save_document(st.session_state.chapter_docs[st.session_state.selected_chapter], st.session_state.selected_chapter)
                else:
                    st.warning("请先生成内容再保存！")

    def save_all_documents(self) -> None:
        """
        保存所有已生成的章节文档
        """
        for chapter_title, doc in st.session_state.chapter_docs.items():
            self.save_document(doc, chapter_title)

    def save_all_in_one_document(self) -> None:
        """
        将所有已生成的章节内容保存在一个文档中
        """
        combined_doc = Document()
        combined_doc.styles[self.doc_styles].font.size = Pt(self.font_size)
        combined_doc.styles[self.doc_styles].font.name = self.font_name
        combined_doc.styles[self.doc_styles].element.rPr.rFonts.set(qn('w:eastAsia'), self.font_name)

        for chapter_title, doc in st.session_state.chapter_docs.items():
            # 添加章节内容
            for element in doc.element.body:
                combined_doc.element.body.append(element)

        file_name = "combined_document.docx"
        combined_doc.save(file_name)
        st.success(f"所有章节内容已保存到 {file_name}！")

        # 添加下载按钮
        with open(file_name, "rb") as file:
            btn = st.download_button(
                label="下载文档",
                data=file,
                file_name=file_name,
                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            )

# 初始化session_state，如果尚未初始化
if 'uploaded_file' not in st.session_state:
    st.session_state.uploaded_file = {}
with st.sidebar:
    st.sidebar.markdown("### 请上传招标文件:", unsafe_allow_html=True)
    uploaded_file = st.sidebar.file_uploader("选择文件", type=["pdf"])
    # 显示上传的文件
    if uploaded_file:
        # 查看st.session_state['uploaded_file']中saved是否为true
        if not st.session_state['uploaded_file'].get('saved', False) or st.session_state['uploaded_file'].get('id') != uploaded_file.file_id:
            # 清空已有上传文件信息，重新处理
            st.session_state.uploaded_file = {}
            # 如果文件已上传，则显示文件信息
            # 创建一个占位符
            placeholder = st.empty()
            placeholder.write('文件正在上传⬆️...')
            # 处理上传文件的代码
            file_details = {"name": uploaded_file.name,
                            "id": uploaded_file.file_id,
                            "type": uploaded_file.type,
                            "saved": False,
                            "project_name": "",
                            "project_no": ""
                            }
            st.session_state['uploaded_file'] = file_details
            # 对pdf文件解析后进行索引
            # project_name, project_no = bid_service.notice_index(uploaded_file)
            project_name, project_no = "test_name", "test_no"
            file_details['saved'] = True
            file_details['project_name'] = project_name
            file_details['project_no'] = project_no
            placeholder.empty()
        st.write("文件上传成功")
        st.write({"项目名称": st.session_state['uploaded_file']['project_name'],
                  "项目编号": st.session_state['uploaded_file']['project_no']})
        # 定义一条分割线
        st.markdown("""
        <style>
        .custom-divider {
            border-top: 3px solid #ff0000;  /* 红色虚线 */
            margin: 10px 0;  /* 上下边距 */
        }
        </style>
        <div class="custom-divider"></div>
        """, unsafe_allow_html=True)
        # 生成目录的逻辑
        # 使用streamlit添加一个用户输入框，用于获取用户的输入
        user_input = st.text_area(label="category needs",label_visibility="hidden")
        _input = ""
        if user_input:
            _input = user_input
        st.write(_input)
        chapter = gen_catalog.chat(_input)
        generate = GenerateChapter(chapter)
        generate.create_launcher()

        # 添加保存所有文档的按钮
        if st.button("保存所有章节文档"):
            generate.save_all_in_one_document()
            st.success("所有章节内容已保存到一个文档中！")

if query:
    log.info('query start...query={}', query)
    # 用户刚刚输入的内容存入messages
    st.session_state['messages'].append({'role': 'human', 'content': query})
    # 界面展示用户刚刚输入的内容
    st.chat_message('human').write(query)
    if 'name' not in st.session_state.uploaded_file:
        st.error('请先上传需要解析的招标文件')
        st.stop()
    if 'saved' not in st.session_state.uploaded_file or not st.session_state.uploaded_file['saved']:
        st.error('请重新上传招标文件')
        st.stop()
    if 'project_name' not in st.session_state.uploaded_file or 'project_no' not in st.session_state.uploaded_file:
        st.error('当前招标文件缺少项目名称/项目编号的关键信息')
        st.stop()
    # 获取project_name project_no
    project_name = st.session_state.uploaded_file['project_name']
    project_no = st.session_state.uploaded_file['project_no']
    # spinner
    with (st.spinner('脑细胞正在燃🔥🔥🔥')):
        # 1.处理文件信息
        file_name = None
        file_bytes = None
        memory = None
        if 'memory' in st.session_state:
            memory: ConversationBufferWindowMemory = st.session_state['memory']
            log.info("memory={}", memory)
        # 1.RAG流程，获取问题的答案
        if not project_name:
            project_name = "梨树县 2025 年高标准农田建设储备项目勘察设计（第二批）1 标段"
        new_query = f"{project_name}{query}"
        response = bid_service.analysis_notice_llm(query)
        log.info('query response={}', response)
        msg = {'role': 'ai', 'content': response}
        # 数据存入st.session_state
        st.session_state['messages'].append(msg)
        # 更新memory
        if memory:
            output_str = {'response': response}
            memory.save_context({"input": query}, {"output": json.dumps(output_str, ensure_ascii=False)})
        # 将数据写出到页面
        st.chat_message('ai').write(response, unsafe_allow_html=True)
